import numpy as np
from .image_convert_base import ConvertBase
import cv2 as cv

class OffsetConvert(ConvertBase):
    
    def __init__(self, use_rate = 0.5, min_offset = -30, max_offset = 30, x_offset = 0, y_offset = 0, is_rand = True):
        super().__init__(use_rate)
        
        self.min_offset = min_offset
        self.max_offset = max_offset
        
        self.x_offset = x_offset
        self.y_offset = y_offset
        
        self.is_rand = is_rand
    
    
    def __get_offset(self):
        return [np.random.randint(self.min_offset, self.max_offset), np.random.randint(self.min_offset, self.max_offset)]
        
        
    def convert(self, img, boxes, points):
        if self.is_rand:
            self.x_offset, self.y_offset = self.__get_offset()
            
        h, w, _ = img.shape
        offset_mat = np.array([[1, 0, self.x_offset], [0, 1, self.y_offset]], np.float32)
        img = cv.warpAffine(img, offset_mat, (w, h))
        
        if boxes is not None:
            boxes[:, [0, 1]] += [self.x_offset, self.y_offset]
            boxes[:, [2, 3]] += [self.x_offset, self.y_offset]
            
        if points is not None:
            points += [self.x_offset, self.y_offset]
        
        return img, boxes, points    
        